{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, xml, xml.etree.ElementTree as ET, numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "*******This is a brief decsription of the i2b2 sample data*********\r\n",
      "1. Data files:\r\n",
      "   The group of 3 files with file extentions: txt, extent and tlink\r\n",
      "are in the format that are consistent with the previous i2b2 challenges.\r\n",
      "   In addition, an xml file is provided for each record. The XML \r\n",
      "file is browsable using the MAE tool found in the directory.\r\n",
      "\r\n",
      "2. Usage of MAE tool\r\n",
      "   - Download and unzip the MAE.zip \r\n",
      "   - Run the MAE_v0.9.3.jar (please make sure that you have Java installed)\r\n",
      "   - In menu bar, select File -> Load DTD, and in the popup window,\r\n",
      "     navigate to the sample folder in the unzipped MAE folder, load the\r\n",
      "     TemporalAnnotation.dtd\r\n",
      "   - Then, select File -> Load File, and load one of the records, e.g. \r\n",
      "     357.xml\r\n",
      "   - You will be able to browse the annotations by clicking on different\r\n",
      "     tabs.\r\n",
      "\r\n"
     ]
    }
   ],
   "source": [
    "cat readme.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "START_CDATA = \"<TEXT><![CDATA[\"\n",
    "END_CDATA   = \"]]></TEXT>\"\n",
    "\n",
    "TAGS        = ['MEDICATION', 'OBSEE', 'SMOKER', 'HYPERTENSION', 'event', 'FAMILY_HIST']\n",
    "\n",
    "def read_xml_file(xml_path, event_tag_type='ALL_CHILDREN', match_text=True):\n",
    "    with open(xml_path, mode='r') as f:\n",
    "        lines = f.readlines()\n",
    "        text, in_text = [], False\n",
    "        for i, l in enumerate(lines):\n",
    "            if START_CDATA in l:\n",
    "                text.append(list(l[l.find(START_CDATA) + len(START_CDATA):]))\n",
    "                in_text = True\n",
    "            elif END_CDATA in l:\n",
    "                text.append(list(l[:l.find(END_CDATA)]))\n",
    "                break\n",
    "            elif in_text:\n",
    "#                 if xml_path.endswith('180-03.xml') and '0808' in l and 'Effingham' in l:\n",
    "#                     print(\"Adjusting known error\")\n",
    "#                     l = l[:9] + ' ' * 4 + l[9:]\n",
    "# #                 elif xml_path.endswith('188-05.xml') and 'Johnson & Johnson' in l:\n",
    "# #                     print(\"Adjusting known error\")\n",
    "# #                     l = l.replace('&', 'and')\n",
    "                text.append(list(l))\n",
    "        \n",
    "    pos_transformer = {}\n",
    "    \n",
    "    linear_pos = 1\n",
    "    for line, sentence in enumerate(text):\n",
    "        for char_pos, char in enumerate(sentence):\n",
    "            pos_transformer[linear_pos] = (line, char_pos)\n",
    "            linear_pos += 1\n",
    "        \n",
    "    try: xml_parsed = ET.parse(xml_path)\n",
    "    except:\n",
    "        print(xml_path)\n",
    "        raise\n",
    "        \n",
    "    tag_containers = xml_parsed.findall('TAGS')\n",
    "    assert len(tag_containers) == 1, \"Found multiple tag sets!\"\n",
    "    tag_container = tag_containers[0]\n",
    "    \n",
    "#     event_tags = tag_container.getchildren() if event_tag_type == 'ALL_CHILDREN' else tag_container.findall('event')\n",
    "    event_tags = tag_container.findall('EVENT')\n",
    "    event_labels = [['O'] * len(sentence) for sentence in text]\n",
    "    for event_tag in event_tags:\n",
    "        base_label = event_tag.attrib['type']\n",
    "        start_pos, end_pos, event_text = event_tag.attrib['start'], event_tag.attrib['end'], event_tag.attrib['text']\n",
    "        start_pos, end_pos = int(start_pos)+1, int(end_pos)\n",
    "        event_text = ' '.join(event_text.split())\n",
    "#         if event_text == \"0808 O’neil’s Court\":\n",
    "#             print(\"Adjusting known error\")\n",
    "#             end_pos -= 4\n",
    "#         if event_text == 'Johnson and Johnson' and xml_path.endswith('188-05.xml'):\n",
    "#             print(\"Adjusting known error\")\n",
    "#             event_text = 'Johnson & Johnson'\n",
    "        \n",
    "\n",
    "        (start_line, start_char), (end_line, end_char) = pos_transformer[start_pos], pos_transformer[end_pos]\n",
    "            \n",
    "        obs_text = []\n",
    "        for line in range(start_line, end_line+1):\n",
    "            t = text[line]\n",
    "            s = start_char if line == start_line else 0\n",
    "            e = end_char if line == end_line else len(t)\n",
    "            obs_text.append(''.join(t[s:e+1]).strip())\n",
    "        obs_text = ' '.join(obs_text)\n",
    "        obs_text = ' '.join(obs_text.split())\n",
    "        \n",
    "        if '&apos;' in obs_text and '&apos;' not in event_text: event_text = event_text.replace(\"'\", \"&apos;\")\n",
    "        if '&quot;' in obs_text and '&quot;' not in event_text: event_text = event_text.replace('\"', '&quot;')\n",
    "              \n",
    "        if match_text: assert obs_text == event_text, (\n",
    "            (\"Texts don't match! %s v %s\" % (event_text, obs_text)) + '\\n' + str((\n",
    "                start_pos, end_pos, line, s, e, t, xml_path\n",
    "            ))\n",
    "        )\n",
    "            \n",
    "        if base_label.strip() == '': continue\n",
    "        \n",
    "        event_labels[end_line][end_char]     = 'I-%s' % base_label\n",
    "        event_labels[start_line][start_char] = 'B-%s' % base_label\n",
    "        \n",
    "        for line in range(start_line, end_line+1):\n",
    "            t = text[line]\n",
    "            s = start_char+1 if line == start_line else 0\n",
    "            e = end_char-1 if line == end_line else len(t)-1\n",
    "            for i in range(s, e+1): event_labels[line][i] = 'I-%s' % base_label\n",
    "\n",
    "    return text, event_labels\n",
    "    \n",
    "def merge_into_words(text_by_char, all_labels_by_char):\n",
    "    assert len(text_by_char) == len(all_labels_by_char), \"Incorrect # of sentences!\"\n",
    "    \n",
    "    N = len(text_by_char)\n",
    "    \n",
    "    text_by_word, all_labels_by_word = [], []\n",
    "    \n",
    "    for sentence_num in range(N):\n",
    "        sentence_by_char = text_by_char[sentence_num]\n",
    "        labels_by_char   = all_labels_by_char[sentence_num]\n",
    "        \n",
    "        assert len(sentence_by_char) == len(labels_by_char), \"Incorrect # of chars in sentence!\"\n",
    "        S = len(sentence_by_char)\n",
    "        \n",
    "        if labels_by_char == (['O'] * len(sentence_by_char)):\n",
    "            sentence_by_word = ''.join(sentence_by_char).split()\n",
    "            labels_by_word   = ['O'] * len(sentence_by_word)\n",
    "        else: \n",
    "            sentence_by_word, labels_by_word = [], []\n",
    "            text_chunks, labels_chunks = [], []\n",
    "            s = 0\n",
    "            for i in range(S):\n",
    "                if i == S-1:\n",
    "                    text_chunks.append(sentence_by_char[s:])\n",
    "                    labels_chunks.append(labels_by_char[s:])\n",
    "                elif labels_by_char[i] == 'O': continue\n",
    "                else:\n",
    "                    if i > 0 and labels_by_char[i-1] == 'O':\n",
    "                        text_chunks.append(sentence_by_char[s:i])\n",
    "                        labels_chunks.append(labels_by_char[s:i])\n",
    "                        s = i\n",
    "                    if labels_by_char[i+1] == 'O' or labels_by_char[i+1][2:] != labels_by_char[i][2:]:\n",
    "                        text_chunks.append(sentence_by_char[s:i+1])\n",
    "                        labels_chunks.append(labels_by_char[s:i+1])\n",
    "                        s = i+1\n",
    "                \n",
    "            for text_chunk, labels_chunk in zip(text_chunks, labels_chunks):\n",
    "                assert len(text_chunk) == len(labels_chunk), \"Bad Chunking (len)\"\n",
    "                assert len(text_chunk) > 0, \"Bad chunking (len 0)\" + str(text_chunks) + str(labels_chunks)\n",
    "                \n",
    "                labels_set = set(labels_chunk)\n",
    "                assert labels_set == set(['O']) or (len(labels_set) <= 3 and 'O' not in labels_set), (\n",
    "                    (\"Bad chunking (contents) %s\" % ', '.join(labels_set))+ str(text_chunks) + str(labels_chunks)\n",
    "                )\n",
    "                \n",
    "                text_chunk_by_word = ''.join(text_chunk).split()\n",
    "                W = len(text_chunk_by_word)\n",
    "                if W == 0: \n",
    "#                     assert labels_set == set(['O']), \"0-word chunking and non-0 label!\" + str(\n",
    "#                         text_chunks) + str(labels_chunks\n",
    "#                     )\n",
    "                    continue\n",
    "                \n",
    "                if labels_chunk[0] == 'O': labels_chunk_by_word = ['O'] * W\n",
    "                elif W == 1:               labels_chunk_by_word = [labels_chunk[0]]\n",
    "                elif W == 2:               labels_chunk_by_word = [labels_chunk[0], labels_chunk[-1]]\n",
    "                else:                      labels_chunk_by_word = [\n",
    "                        labels_chunk[0]\n",
    "                    ] + [labels_chunk[1]] * (W - 2) + [\n",
    "                        labels_chunk[-1]\n",
    "                    ]\n",
    "                    \n",
    "                sentence_by_word.extend(text_chunk_by_word)\n",
    "                labels_by_word.extend(labels_chunk_by_word)\n",
    "\n",
    "        assert len(sentence_by_word) == len(labels_by_word), \"Incorrect # of words in sentence!\"    \n",
    "        \n",
    "        if len(sentence_by_word) == 0: continue\n",
    "            \n",
    "        text_by_word.append(sentence_by_word)\n",
    "        all_labels_by_word.append(labels_by_word)\n",
    "    return text_by_word, all_labels_by_word\n",
    "\n",
    "def reprocess_event_labels(folders, base_path='.', event_tag_type='event', match_text=True, dev_set_size=None):\n",
    "    all_texts_by_patient, all_labels_by_patient = {}, {}\n",
    "\n",
    "    for folder in folders:\n",
    "        folder_dir = os.path.join(base_path, folder)\n",
    "        xml_filenames = [x for x in os.listdir(folder_dir) if x.endswith('xml')]\n",
    "        for xml_filename in xml_filenames:\n",
    "            patient_num = int(xml_filename[:-4])\n",
    "            xml_filepath = os.path.join(folder_dir, xml_filename)\n",
    "            \n",
    "            text_by_char, labels_by_char = read_xml_file(\n",
    "                xml_filepath,\n",
    "                event_tag_type=event_tag_type,\n",
    "                match_text=match_text\n",
    "            )\n",
    "            text_by_word, labels_by_word = merge_into_words(text_by_char, labels_by_char)\n",
    "            \n",
    "            if patient_num not in all_texts_by_patient:\n",
    "                all_texts_by_patient[patient_num] = []\n",
    "                all_labels_by_patient[patient_num] = []\n",
    "            \n",
    "            all_texts_by_patient[patient_num].extend(text_by_word)\n",
    "            all_labels_by_patient[patient_num].extend(labels_by_word)\n",
    "            \n",
    "    patients = set(all_texts_by_patient.keys())\n",
    "    \n",
    "    if dev_set_size is None: train_patients, dev_patients = list(patients), []\n",
    "    else:\n",
    "        N_train = int(len(patients) * (1-dev_set_size))\n",
    "        patients_random = np.random.permutation(list(patients))\n",
    "        train_patients = list(patients_random[:N_train])\n",
    "        dev_patients   = list(patients_random[N_train:])\n",
    "    \n",
    "    train_texts, train_labels = [], []\n",
    "    dev_texts, dev_labels = [], []\n",
    "    \n",
    "    for patient_num in train_patients:\n",
    "        train_texts.extend(all_texts_by_patient[patient_num])\n",
    "        train_labels.extend(all_labels_by_patient[patient_num])\n",
    "\n",
    "    for patient_num in dev_patients:\n",
    "        dev_texts.extend(all_texts_by_patient[patient_num])\n",
    "        dev_labels.extend(all_labels_by_patient[patient_num])\n",
    "\n",
    "\n",
    "    train_out_text_by_sentence = []\n",
    "    for text, labels in zip(train_texts, train_labels):\n",
    "        train_out_text_by_sentence.append('\\n'.join('%s %s' % x for x in zip(text, labels)))\n",
    "    dev_out_text_by_sentence = []\n",
    "    for text, labels in zip(dev_texts, dev_labels):\n",
    "        dev_out_text_by_sentence.append('\\n'.join('%s %s' % x for x in zip(text, labels)))\n",
    "\n",
    "    return '\\n\\n'.join(train_out_text_by_sentence), '\\n\\n'.join(dev_out_text_by_sentence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_train_text, final_dev_text = reprocess_event_labels(\n",
    "    ['2012-07-15.original-annotation.release'], dev_set_size=0.1, match_text=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_text, _ = reprocess_event_labels(\n",
    "    ['2012-08-08.test-data.event-timex-groundtruth/xml'], match_text=False, dev_set_size=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ADMISSION B-OCCURRENCE\n",
      "DATE O\n",
      ": O\n",
      "\n",
      "12/13/2002 O\n",
      "\n",
      "DISCHARGE B-OCCURRENCE\n",
      "DATE O\n",
      ": O\n",
      "\n",
      "12/14/2002 O\n",
      "\n",
      "HISTORY O\n",
      "OF O\n",
      "PRESENT O\n",
      "ILLNESS O\n",
      ": O\n",
      "\n",
      "The O\n",
      "patient O\n",
      "is O\n",
      "a O\n",
      "30 O\n",
      "-year-old O\n",
      "Gravida B-OCCURRENCE\n",
      "III O\n",
      ", O\n",
      "Para B-OCCURRENCE\n",
      "II O\n",
      "who O\n",
      "presented B-EVIDENTIAL\n",
      "desiring O\n",
      "definitive B-TREATMENT\n",
      "surgical I-TREATMENT\n",
      "sterilization I-TREATMENT\n",
      ". O\n",
      "\n",
      "She O\n",
      "was O\n",
      "extensively B-OCCURRENCE\n",
      "counseled I-OCCURRENCE\n",
      "regarding O\n",
      "other B-TREATMENT\n",
      "methods I-TREATMENT\n",
      "of O\n",
      "birth B-TREATMENT\n",
      "control I-TREATMEN\n"
     ]
    }
   ],
   "source": [
    "print(final_train_text[:500])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Admission B-OCCURRENCE\n",
      "Date O\n",
      ": O\n",
      "\n",
      "2010-02-05 O\n",
      "\n",
      "Discharge B-OCCURRENCE\n",
      "Date O\n",
      ": O\n",
      "\n",
      "2010-02-06 O\n",
      "\n",
      "Service O\n",
      ": O\n",
      "\n",
      "TRAUMA O\n",
      "\n",
      "HISTORY O\n",
      "OF O\n",
      "PRESENT O\n",
      "ILLNESS O\n",
      ": O\n",
      "\n",
      "The O\n",
      "patient O\n",
      "is O\n",
      "a O\n",
      "23 O\n",
      "year O\n",
      "old O\n",
      "female O\n",
      ", O\n",
      "status O\n",
      "post O\n",
      "fall B-PROBLEM\n",
      "from O\n",
      "standing O\n",
      "position O\n",
      "after O\n",
      "slipping O\n",
      "on O\n",
      "ice O\n",
      ". O\n",
      "\n",
      "She O\n",
      "had O\n",
      "no O\n",
      "loss B-PROBLEM\n",
      "of I-PROBLEM\n",
      "consciousness I-PROBLEM\n",
      ". O\n",
      "\n",
      "She O\n",
      "recall\n"
     ]
    }
   ],
   "source": [
    "print(final_dev_text[:400])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ADMISSION B-OCCURRENCE\n",
      "DATE O\n",
      ": O\n",
      "\n",
      "10/14/96 O\n",
      "\n",
      "DISCHARGE B-OCCURRENCE\n",
      "DATE O\n",
      ": O\n",
      "\n",
      "10/27/96 O\n",
      "date O\n",
      "of O\n",
      "birth B-OCCURRENCE\n",
      "; O\n",
      "September O\n",
      "30 O\n",
      ", O\n",
      "1917 O\n",
      "\n",
      "THER O\n",
      "PROCEDURES O\n",
      ": O\n",
      "\n",
      "arterial B-TEST\n",
      "catheterization I-TEST\n",
      "on O\n",
      "10/14/96 O\n",
      ", O\n",
      "head B-TEST\n",
      "CT I-TEST\n",
      "scan I-TEST\n",
      "on O\n",
      "10/14/96 O\n",
      "\n",
      "HISTORY O\n",
      "AND O\n",
      "REASON O\n",
      "FOR O\n",
      "HOSPITALIZATION O\n",
      ": O\n",
      "\n",
      "Granrivern O\n",
      "Call O\n",
      "is O\n",
      "a O\n",
      "79-year-old O\n",
      "right O\n",
      "han\n"
     ]
    }
   ],
   "source": [
    "print(test_text[:400])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = {}\n",
    "for s in final_train_text, final_dev_text, test_text:\n",
    "    for line in s.split('\\n'):\n",
    "        if line == '': continue\n",
    "        label = line.split()[-1]\n",
    "        assert label == 'O' or label.startswith('B-') or label.startswith('I-'), \"label wrong! %s\" % label\n",
    "        if label not in labels: labels[label] = 1\n",
    "        else: labels[label] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'B-OCCURRENCE': 5774,\n",
       " 'O': 114910,\n",
       " 'B-EVIDENTIAL': 1334,\n",
       " 'B-TREATMENT': 7098,\n",
       " 'I-TREATMENT': 6748,\n",
       " 'I-OCCURRENCE': 3590,\n",
       " 'B-CLINICAL_DEPT': 1724,\n",
       " 'I-CLINICAL_DEPT': 3253,\n",
       " 'B-PROBLEM': 9319,\n",
       " 'I-PROBLEM': 13543,\n",
       " 'B-TEST': 4762,\n",
       " 'I-TEST': 5931,\n",
       " 'I-EVIDENTIAL': 84}"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./processed/train.tsv', mode='w') as f:\n",
    "    f.write(final_train_text)\n",
    "with open('./processed/dev.tsv', mode='w') as f:\n",
    "    f.write(final_dev_text)\n",
    "with open('./processed/test.tsv', mode='w') as f:\n",
    "    f.write(test_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
